﻿/*
* ////////////////////////////////////////////////
*
* 开发者：AZPEP
* 友情链接：
*	-	B站：		零分的考卷			https://space.bilibili.com/257391258
*	-	Github：	AZPEP-BiliBili		https://github.com/AZPEP-BiliBili
*	-	Gitee：		AZPEP_BiliBili		https://gitee.com/AZPEP_BiliBili
*
* CopyRight AZPEP
*
* ////////////////////////////////////////////////
*/

#include "pch.h"
#include "GetDrive.h"

/// <summary>
/// 获取所有驱动器的信息列表
/// </summary>
/// <param name="DriveInformationList">指向DRIVE_LAYOUT_INFORMATION_EX*的列表</param>
/// <returns>成功返回0，DriveInformationList为指向DRIVE_LAYOUT_INFORMATION_EX*的列表；失败返回GetLastError()得到的错误代码，同时DriveInformationList为空列表</returns>
DWORD MountEsp_GetAllDriveLayoutEx(
	_Out_ std::vector<DRIVE_LAYOUT_INFORMATION_EX*>* DriveInformationList
)
{
	/* 获取磁盘总数 */
	DWORD dDiskNumber;											//磁盘总数
	DWORD dStatus = MountEsp_GetDiskCount(&dDiskNumber);

	if (dStatus != 0) return dStatus;							//获取失败，返回错误代码

	/* 获取DRIVE_LAYOUT_INFORMATION_EX */
	std::vector<DRIVE_LAYOUT_INFORMATION_EX*> tmp_vec_DLIE;		//临时存放DRIVE_LAYOUT_INFORMATION_EX*的列表

	//通过遍历获取DRIVE_LAYOUT_INFORMATION_EX
	for (DWORD d = 0; d < dDiskNumber; d++) {
		MallocEx(DRIVE_LAYOUT_INFORMATION_EX*, _DLIE, MAX_DRIVE_LAYOUT_INFORMATION_EX);

		//打开磁盘
		std::wstringstream wss_PhysicalDrivePath;
		wss_PhysicalDrivePath << L"\\\\.\\PhysicalDrive";
		wss_PhysicalDrivePath << d;

		HANDLE hDevice = CreateFile(
			wss_PhysicalDrivePath.str().c_str(),
			GENERIC_READ,
			FILE_SHARE_READ | FILE_SHARE_WRITE,
			NULL,
			OPEN_EXISTING,
			0,
			NULL
		);

		//打开失败，返回错误代码
		if (hDevice == INVALID_HANDLE_VALUE) { DriveInformationList->clear(); return GetLastError(); }

		//获取DRIVE_LAYOUT_INFORMATION_EX
		BOOL bStatus = DeviceIoControl(
			hDevice,
			IOCTL_DISK_GET_DRIVE_LAYOUT_EX,
			NULL,
			0,
			_DLIE,
			MAX_DRIVE_LAYOUT_INFORMATION_EX + MAX_PATH,
			NULL,
			NULL
		);

		//获取失败，返回错误代码
		if (bStatus != TRUE) { DriveInformationList->clear(); return GetLastError(); }

		//关闭句柄
		CloseHandle(hDevice);

		//保存结构体
		tmp_vec_DLIE.push_back(_DLIE);
	}

	for (int i = 0; i < tmp_vec_DLIE.size(); i++) {
		DriveInformationList->push_back(tmp_vec_DLIE[i]);
	}
	//free(&tmp_vec_DLIE);
	return 0;
}

/// <summary>
/// 获取计算机上装载的磁盘数
/// </summary>
/// <param name="dDiskNumber">计算机上装载的磁盘数</param>
/// <returns>成功返回0，且dDiskNumber等于装载的磁盘数；失败返回GetLastError()得到的错误代码，且dDiskNumber的值为-1</returns>
DWORD MountEsp_GetDiskCount(
	_Out_ DWORD* dDiskNumber
) {
	*dDiskNumber = 0;
	while (true) {
		// 拼接硬盘路径，拼接后路径 \\\\.\\PhysicalDriveX (X是硬盘号)
		std::wstringstream wss_PhysicalDrivePath;
		wss_PhysicalDrivePath << L"\\\\.\\PhysicalDrive";
		wss_PhysicalDrivePath << *dDiskNumber;

		HANDLE hDevice = INVALID_HANDLE_VALUE; //硬盘句柄
		//打开磁盘
		hDevice = CreateFile(
			wss_PhysicalDrivePath.str().c_str(),
			GENERIC_READ,
			FILE_SHARE_READ | FILE_SHARE_WRITE,
			NULL,
			OPEN_EXISTING,
			0,
			NULL
		);

		if (hDevice == INVALID_HANDLE_VALUE) {
			DWORD ErrorCode = GetLastError();
			//没有磁盘，跳出循环
			if (ErrorCode == ERROR_FILE_NOT_FOUND) break;

			//读取失败，报错返回
			else {
				*dDiskNumber = -1;
				return ErrorCode;
			}
		}

		//否则，磁盘号递增
		CloseHandle(hDevice);
		*dDiskNumber = *dDiskNumber + 1;
	}

	return 0;
}